#include <iostream>
#include "TFile.h"
#include "TTree.h"
#include "TH1D.h"
#include <cmath>
#include <algorithm>
#include <map>
#include <sstream>
#include "evaluate.h"
#include "SimpleTrack.h"


using namespace std;


class floatBin
{
  public:
    floatBin(float l, float h) : low(l), high(h) {}
    ~floatBin(){}
    float low;
    float high;
    bool operator <(const floatBin& other) const
    {
      return ( high < other.low );
    }
};


//return which bin in bins that pt belongs to, or -1 if it doesn't belong to any
static int pt_bin(vector<floatBin>& bins, float pt)
{
  floatBin bin(pt,pt);
  if( ( bin < bins[0] ) || ( bins.back() < bin ) ){return -1;}
  pair<vector<floatBin>::iterator,vector<floatBin>::iterator> bounds;
  bounds = equal_range(bins.begin(), bins.end(), bin);
  return ( (int)(bounds.first - bins.begin()) );
}


// pT in GeV/c , B in T , kappa in cm^-1
static inline double pT_to_kappa(double B, double pT)
{
  return 0.003*B/pT;
}


static inline double kappa_to_pT(double B, double kappa)
{
  return 0.003*B/kappa;
}


// calculate_eff <input file> <histogram output file>
int main(int argc, char** argv)
{
  //define pt bins to use
  float xbins[27];
  vector<floatBin> pt_bins;
  for(int i=2;i<=9;++i)
  {
    pt_bins.push_back( floatBin( ((float)(i))*0.1, ((float)(i+1))*0.1 ) );
    xbins[i-2] = ((float)(i))*0.1;
  }
  float min_bin = 1.0;
  for(int i=8;i<=25;++i)
  {
    pt_bins.push_back(floatBin(min_bin, min_bin+0.5));xbins[i] = min_bin;
    min_bin += 0.5;
  }
  xbins[26] = 10.0;
  
  TH1D n_mc("n_mc", "# mc tracks vs pt", pt_bins.size(), xbins);
  TH1D n_reco("n_reco", "# reco tracks vs pt", pt_bins.size(), xbins);
  TH1D n_fake("n_fake", "# fake tracks vs pt", pt_bins.size(), xbins);
  TH1D efficiency("efficiency", "efficiency vs pt", pt_bins.size(), xbins);
  TH1D fake_rate("fake_rate", "fake rate vs pt", pt_bins.size(), xbins);
  TH1D mom_res("mom_res", "momentum resolution vs pt", pt_bins.size(), xbins);
  
  stringstream ss;
  
  vector<TH1D*> momresbins;
  for(unsigned int i=0;i<pt_bins.size();++i)
  {
    ss.clear();ss.str("");
    ss<<"momresbin_"<<i;
    momresbins.push_back( new TH1D(ss.str().c_str(), ss.str().c_str(), 1024, -1., 1.) );
  }
  
  TH1D purity_4("purity_4", "proportion of tracks with 4 correct hits", pt_bins.size(), xbins);
  TH1D purity_5("purity_5", "proportion of tracks with 4 or 5 correct hits", pt_bins.size(), xbins);
  TH1D purity_6("purity_6", "proportion of tracks with 4, 5, or 6 correct hits", pt_bins.size(), xbins);
  vector<unsigned int> n_4;n_4.assign(pt_bins.size(), 0);
  vector<unsigned int> n_5;n_5.assign(pt_bins.size(), 0);
  vector<unsigned int> n_6;n_6.assign(pt_bins.size(), 0);
  
  TFile mc_file(argv[1]);
  TTree* mc_tree = 0;
  mc_file.GetObject("mc_events", mc_tree);
  SimpleMCEvent* mc_event=0;
  mc_tree->SetBranchAddress("tracks", &mc_event);
  
  TFile reco_file(argv[1]);
  TTree* reco_tree = 0;
  reco_file.GetObject("reco_events", reco_tree);
  SimpleRecoEvent* reco_event=0;
  reco_tree->SetBranchAddress("tracks", &reco_event);
  
  vector<vector<vector<unsigned int> > > mctracks;//pt bin, mc track index, hit index
  mctracks.assign(pt_bins.size(), vector<vector<unsigned int> >());
  vector<vector<float> > mc_mom;
  mc_mom.assign(pt_bins.size(), vector<float>());
  
  vector<vector<unsigned int> > recotracks;//reco track index, hit index
  vector<float> reco_mom;
  
  for(int ev=0;ev<mc_tree->GetEntries();++ev)
  {
    mc_file.cd();
    mc_tree->GetEntry(ev);
    
    for(unsigned int i=0;i<mctracks.size();++i)
    {
      mctracks[i].clear();
      mc_mom[i].clear();
    }
    for(unsigned int i=0;i<mc_event->tracks.size();++i)
    {
      float pt = kappa_to_pT(2.0, mc_event->tracks[i].kappa);
      int bin = pt_bin(pt_bins, pt);
      if(bin < 0){continue;}
      mctracks[bin].push_back(vector<unsigned int>());
      mc_mom[bin].push_back(kappa_to_pT(2.0, mc_event->tracks[i].kappa));
      for(unsigned int h=0;h<mc_event->tracks[i].hits.size();++h)
      {
        mctracks[bin].back().push_back(mc_event->tracks[i].hits[h].index);
      }
    }
    
    reco_file.cd();
    reco_tree->GetEntry(ev);
    
    recotracks.clear();
    reco_mom.clear();
    
    for(unsigned int i=0;i<reco_event->tracks.size();++i)
    {
      recotracks.push_back(vector<unsigned int>());
      reco_mom.push_back(kappa_to_pT(2.0, reco_event->tracks[i].kappa));
      for(unsigned int h=0;h<reco_event->tracks[i].indexes.size();++h)
      {
        recotracks.back().push_back(reco_event->tracks[i].indexes[h]);
      }
    }
    
    vector<bool> mc_is_reconstructed;
    vector<bool> reco_used;reco_used.assign(recotracks.size(), false);
    
    for(unsigned int p=0;p<pt_bins.size();++p)
    {
      map<unsigned int, unsigned int> contr;
      unsigned int nreco = 0;
      calculateEfficiency(mctracks[p], recotracks, 4, nreco, mc_is_reconstructed, reco_used, mc_mom[p], reco_mom, *(momresbins[p]), contr);
      
      map<unsigned int, unsigned int>::iterator it;
      it = contr.find(4);
      if(it != contr.end()){n_4[p] += it->second;}
      it = contr.find(5);
      if(it != contr.end()){n_5[p] += it->second;}
      it = contr.find(6);
      if(it != contr.end()){n_6[p] += it->second;}
      
      n_mc.SetBinContent(p+1, n_mc.GetBinContent(p+1) + mctracks[p].size());
      n_reco.SetBinContent(p+1, n_reco.GetBinContent(p+1) + nreco);
    }
    for(unsigned int i=0;i<reco_used.size();++i)
    {
      if(reco_used[i] == false)
      {
        float pt = kappa_to_pT(2.0, reco_event->tracks[i].kappa);
        int bin = pt_bin(pt_bins, pt);
        if(bin < 0){continue;}
        n_fake.SetBinContent(bin+1, n_fake.GetBinContent(bin+1) + 1);
      }
    }
  }
  
  for(unsigned int i=0;i<pt_bins.size();++i)
  {
    mom_res.SetBinContent(i+1, momresbins[i]->GetRMS());
  }
  
  n_mc.Sumw2();
  n_reco.Sumw2();
  n_fake.Sumw2();
  
  efficiency.Divide(&n_reco, &n_mc, 1., 1., "B");
  fake_rate.Divide(&n_fake, &n_mc);
  
  for(unsigned int p=0;p<pt_bins.size();++p)
  {
    cout<<p<<" "<<n_4[p]<<" "<<n_5[p]<<" "<<n_6[p]<<endl;
    unsigned int ntot = n_4[p] + n_5[p] + n_6[p];
    if(ntot != 0)
    {
      purity_4.SetBinContent(p+1, ((float)(n_4[p]))/((float)(ntot)));
      purity_5.SetBinContent(p+1, ((float)(n_4[p] + n_5[p]))/((float)(ntot)));
      purity_6.SetBinContent(p+1, 1.);
    }
    else
    {
      purity_4.SetBinContent(p+1, 0.);
      purity_5.SetBinContent(p+1, 0.);
      purity_6.SetBinContent(p+1, 0.);
    }
  }
  
  TFile out(argv[2], "recreate");
  out.cd();
  n_mc.Write();
  n_reco.Write();
  n_fake.Write();
  efficiency.Write();
  fake_rate.Write();
  mom_res.Write();
  purity_4.Write();
  purity_5.Write();
  purity_6.Write();
  
  return 0;
}


